from vnpy.app.cta_strategy import (
    CtaTemplate,
    StopOrder,
    TickData,
    BarData,
    TradeData,
    OrderData,
    BarGenerator,
    ArrayManager
)
from vnpy.trader.constant import Interval, Direction, Offset


class BTCLeveragedNeutralGridStrategy(CtaTemplate):
    author = "Your Name"

    # 策略参数
    grid_step = 100  # 网格步长
    price_range_low = 30000  # 价格区间下限
    price_range_high = 40000  # 价格区间上限
    init_volume = 1  # 初始开仓数量

    parameters = ["grid_step", "price_range_low", "price_range_high", "init_volume"]
    variables = []

    def __init__(self, cta_engine, strategy_name, vt_symbol, setting):
        """"""
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)

        self.bg = BarGenerator(self.on_bar, 1, self.on_1min_bar, Interval.MINUTE)
        self.am = ArrayManager()

    def on_init(self):
        """
        Callback when strategy is inited.
        """
        self.write_log("策略初始化")
        self.load_bar(10)

    def on_start(self):
        """
        Callback when strategy is started.
        """
        self.write_log("策略启动")

    def on_stop(self):
        """
        Callback when strategy is stopped.
        """
        self.write_log("策略停止")
        self.cancel_all()

    def on_tick(self, tick: TickData):
        """
        Callback of new tick data update.
        """
        self.bg.update_tick(tick)

    def on_bar(self, bar: BarData):
        """
        Callback of new bar data update.
        """
        self.bg.update_bar(bar)

    def on_1min_bar(self, bar: BarData):
        """
        Callback of new 1-minute bar data update.
        """
        self.am.update_bar(bar)
        if not self.am.inited:
            return

        current_price = bar.close_price

        # 检查价格是否在区间内
        if self.price_range_low <= current_price <= self.price_range_high:
            # 计算当前网格位置
            grid_level = int((current_price - self.price_range_low) / self.grid_step)

            # 开多仓条件
            if current_price <= self.price_range_low + grid_level * self.grid_step:
                self.buy(current_price, self.init_volume)

            # 平多仓条件
            if current_price >= self.price_range_low + (grid_level + 1) * self.grid_step:
                self.sell(current_price, self.init_volume)

        self.put_orders()

    def on_trade(self, trade: TradeData):
        """
        Callback of new trade data update.
        """
        self.put_orders()

    def on_order(self, order: OrderData):
        """
        Callback of new order data update.
        """
        pass

    def put_orders(self):
        """
        Place orders based on strategy logic.
        """
        pass
